
#### Empty matrix w/ indexes ----
forCycle <- function(x){
  # create an empty vector
  v <- c()
  for(i in 1:371){
    v[i] <- i }
  return(v) }


#### Better logs (no 0+1 for zero values) ----
inv_hyp <- function(y){  
  re <- log(y+(((y^2)+1)^(1/2)))
  return(re)}


#### Clustered standard errors ----
cl <- function(dat,fm,cluster) { 
  ef <- estfun(fm)
  if (length(cluster)!=nrow(ef)) {
    cluster <- cluster[as.numeric(rownames(ef))]
  }
  M <- length(unique(cluster))
  N <- length(cluster)
  dfc <- (M/(M-1))*((N-1)/(N-fm$rank))
  u <- apply(estfun(fm),2,
             function(x) tapply(x, cluster, sum))
  vcovCL <- dfc*sandwich(fm, meat=crossprod(u)/N)
  coeftest(fm, vcovCL)
}


#### Standardized coefficients ----
beta.z <- function (MOD) {
  b <- summary(MOD)$coef[-1, 1]
  sx <- sapply(MOD$model[-1], sd)
  sy <- sapply(MOD$model[1], sd)
  beta <- b * sx/sy
  return(beta)
}


#### Conditinal effects for certain values of the moderator ----
conditional_effects <-  function(model, x, m, df, cl, quantiles = 10){
  interact = paste0(x,':',m)
  beta.hat = coef(model) 
  covs = vcovCL(model, df$cl)
  z0  = quantile(model$model[,m], seq(0 , 1, 1/quantiles))
  dy.dx = beta.hat[x] + beta.hat[interact]*z0
  se.dy.dx = sqrt(covs[x, x] + z0^2*covs[interact, interact] + 2*z0*covs[x, interact])
  upr = dy.dx+1.96*se.dy.dx
  lwr = dy.dx-1.96*se.dy.dx
  data.frame(m = z0, b = dy.dx, lwr, upr)
}

#### Function 1: Implementing the Johnson???Neyman Technique ----
jnt <- function(.lm, predictor, moderator, alpha=.05) {
  require(stringi)
  b1 = coef(.lm)[predictor]
  b3 = coef(.lm)[stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor))]
  se_b1 = coef(summary(.lm))[predictor, 2]
  se_b3 = coef(summary(.lm))[stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor)), 2]
  COV_b1b3 = vcov(.lm)[predictor, stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor))]
  t_crit = qt(1-alpha/2, .lm$df.residual)
  # see Bauer & Curran, 2005
  a = t_crit^2 * se_b3^2 - b3^2
  b = 2 * (t_crit^2 * COV_b1b3 - b1 * b3)
  c = t_crit^2 * se_b1^2 - b1^2
  jn = c(
    (-b - sqrt(b^2 - 4 * a * c)) / (2 * a),
    (-b + sqrt(b^2 - 4 * a * c)) / (2 * a)
  )
  JN = sort(unname(jn))
  JN = JN[JN>=min(.lm$model[,moderator]) & JN<=max(.lm$model[,moderator])]
  JN
}

#### Function 2: Plotting theta as a function of Z
theta_plot <- function(.lm, predictor, moderator, alpha=.05, jn=F) {
  require(dplyr)
  require(ggplot2); theme_set(theme_minimal())
  require(stringi)
  .data = data_frame(b1 = coef(.lm)[predictor],
                     b3 = coef(.lm)[stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor))],
                     Z = quantile(.lm$model[,moderator], seq(0,1,.01)),
                     theta = b1 + Z * b3,
                     se_b1 = coef(summary(.lm))[predictor, 2],
                     COV_b1b3 = vcov(.lm)[predictor, stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor))],
                     se_b3 = coef(summary(.lm))[stri_startswith_fixed(names(coef(.lm)), paste0(predictor,":")) | stri_endswith_fixed(names(coef(.lm)), paste0(":",predictor)), 2],
                     se_theta = sqrt(se_b1^2 + 2 * Z * COV_b1b3 + Z^2 * se_b3^2),
                     ci.lo_theta = theta+qt(alpha/2, .lm$df.residual)*se_theta,
                     ci.hi_theta = theta+qt(1-alpha/2, .lm$df.residual)*se_theta)
  if (jn) {
    JN = jnt(.lm=.lm, predictor=predictor, moderator=moderator, alpha=alpha)
    JN_lines = geom_vline(xintercept=JN, linetype=2)
    JN_regions = ifelse(length(JN) == 0, "no significance regions", paste(round(JN,2), collapse = "; "))
    Xlab = paste0(moderator, " (JN Significance Regions: ", JN_regions,")")
  }
  else {
    Xlab = moderator
    JN_lines = NULL
  }
  .data %>%
    ggplot(aes(Z, theta, ymin=ci.lo_theta, ymax=ci.hi_theta)) + geom_ribbon(alpha = .2) + geom_line() + ggtitle(paste("Conditional Effect of", predictor, "as function of", moderator)) + geom_hline(yintercept=0, linetype=2) + labs(x = Xlab, y=expression(theta)) + JN_lines
}

#### Robust standad error for logistic ----
robustse <- function(x, coef = c("logit", "odd.ratio", "probs")) {
  suppressMessages(suppressWarnings(library(lmtest)))
  suppressMessages(suppressWarnings(library(sandwich)))
  
  sandwich1 <- function(object, ...) sandwich(object) *
    nobs(object) / (nobs(object) - 1)
  # Function calculates SE's
  mod1 <- coeftest(x, vcov = sandwich1) 
  # apply the function over the variance-covariance matrix
  
  if (coef == "logit") {
    return(mod1) # return logit with robust SE's
  } else if (coef == "odd.ratio") {
    mod1[, 1] <- exp(mod1[, 1]) # return odd ratios with robust SE's
    mod1[, 2] <- mod1[, 1] * mod1[, 2]
    return(mod1)
  } else {
    mod1[, 1] <- (mod1[, 1]/4) # return probabilites with robust SE's
    mod1[, 2] <- mod1[, 2]/4
    return(mod1)
  }
}


#### Extract RDD ----
extractrdd <- function(dat){
  return(as.numeric(c(dat$coef[1],
                      dat$ci[3,],
                      dat$pv[3,],
                      dat$bws[1],
                      dat$N_h_l,
                      dat$N_h_r)))
}


#### Sandwich standard errors for probit (matches with Stata up to four digits)
#sandwich.rse <- function(object, ...) sandwich(object) * nobs(object) / (nobs(object) - 1)

interaction_plot_threeway_ylim <- function(model, effect, moderator1, moderator2, varcov="default", minimum="min", maximum="max", incr="default", num_points = 50, conf=.95, mean=FALSE, median=FALSE, alph=80, rugplot=T, histogram=T, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient"){
  
  #num_points = 50; conf=.95; alph=80; title="Marginal effects plot"; xlabel="Value of moderator"; ylabel="Estimated marginal coefficient"; minimum="min"; maximum="max"; incr="default"
  #model <- m5
  #effect="ln_price" 
  #moderator1="ln_suitability"
  #moderator2="orgs_locales_median"
  #varcov = covMat5
  
  x <- effect
  z <- moderator1
  w <- moderator2
  xz <- paste0(z, ":", x) 
  xw <- paste0(x, ":", w)
  zw <- paste0(z, ":", w)
  xzw <- paste0(z, ":", x, ":", w)
  
  
  # Define a function to make colors transparent
  makeTransparent<-function(someColor, alpha=alph){
    newColor<-col2rgb(someColor)
    apply(newColor, 2, function(curcoldata){rgb(red=curcoldata[1], green=curcoldata[2],
                                                blue=curcoldata[3],alpha=alpha, maxColorValue=255)})
  }
  
  
  # Extract Variance Covariance matrix

    covMat = varcov
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[x]]
  beta_4 = model$coefficients[[xz]]
  beta_5 = model$coefficients[[xw]]
  beta_7 = model$coefficients[[xzw]]
  
  # Set range of the moderator variable
  # Minimum
  if (minimum == "min"){
    min_val = min(mod_frame[[z]])
  }else{
    min_val = minimum
  }
  # Maximum
  if (maximum == "max"){
    max_val = max(mod_frame[[z]])
  }else{
    max_val = maximum
  }
  
  # Check if minimum smaller than maximum
  if (min_val > max_val){
    stop("Error: Minimum moderator value greater than maximum value.")
  }
  
  # Determine intervals between values of the moderator
  if (incr == "default"){
    increment = (max_val - min_val)/(num_points - 1)
  }else{c
    increment = incr
  }
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- seq(from=min_val, to=max_val, by=increment)
  
  # Compute both marginal effects for 
  delta_1 = beta_1 + beta_4*x_2
  delta_2 = beta_1 + beta_4*x_2 + beta_5 + beta_7*x_2
  
  # Compute variances
  var_1 = covMat[x,x] + 
    (x_2^2)*covMat[xz, xz] + 
    2*x_2*covMat[x, xz] 
  
  var_2 = covMat[x,x] + 
    (x_2^2)*covMat[xz, xz] + 
    (1^2)*covMat[xw, xw] + 
    (x_2^2)*(1^2)*covMat[xzw, xzw] + 
    2*x_2*covMat[x, xz] + 
    2*1*covMat[x, xw] + 
    2*x_2*1*covMat[x, xzw] + 
    2*x_2*1*covMat[xz, xw] + 
    2*(x_2^2)*1*covMat[xz, xzw] + 
    2*x_2*(1^2)*covMat[xw, xzw]  
  
  # Standard errors
  se_1 = sqrt(var_1)
  se_2 = sqrt(var_2)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  
  upper_bound_1 = delta_1 + z_score*se_1
  lower_bound_1 = delta_1 - z_score*se_1
  
  upper_bound_2 = delta_2 + z_score*se_2
  lower_bound_2 = delta_2 - z_score*se_2
  
  # Determine the bounds of the graphing area
  max_y = max(unlist(list(upper_bound_1,upper_bound_2)))
  min_y = min(unlist(list(lower_bound_1,lower_bound_2)))
  
  # Make the histogram color
  alph=80
  hist_col = makeTransparent("grey")
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(-.25, 1), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title, axes=F)
  axis(side=1)
  axis(side=2)
  
  # Plot estimated effects
  lines(y=delta_1, x=x_2, col="black", lty =2, lwd = 2)
  lines(y=upper_bound_1, x=x_2, lty=2)
  lines(y=lower_bound_1, x=x_2, lty=2)
  
  polygon(x = c(x_2, rev(x_2)),
          y = c(lower_bound_1, 
                rev(upper_bound_1)),
          col =  adjustcolor("blue", alpha.f = 0.30), border = NA)
  
  lines(y=delta_2, x=x_2, col="darkgrey", lty = 1, lwd = 2)
  lines(y=upper_bound_2, x=x_2, lty=1)
  lines(y=lower_bound_2, x=x_2, lty=1)
  
  polygon(x = c(x_2, rev(x_2)),
          y = c(lower_bound_2, 
                rev(upper_bound_2)),
          col =  adjustcolor("red", alpha.f = 0.30), border = NA)
  
  
  legend(x=max_val-max_val/10,y=max_y-max_y/10, c("Low", "High"), 
         fill = c(adjustcolor("blue", alpha.f = 0.30), adjustcolor("red", alpha.f = 0.30)), cex=1, seg.len=2,
         text.width=.6, bty = "n")
  #legend(x=max_val-max_val/10,y=max_y-max_y/10,c("Low", "High"), 
  #lty=c(2,1), col = c("black","darkgrey"), lwd = 2, cex=.75, seg.len=2,
  #text.width=.6, bty = "n")
  #c(expression(""<= "Median"), expression(""> "Median"))
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  # Add a vertical line at the mean
  if (mean){
    abline(v = mean(mod_frame[[z]]), lty=2, col="red")
  }
  
  # Add a vertical line at the median
  if (median){
    abline(v = median(mod_frame[[z]]), lty=3, col="blue")
  }
  
  # Add Rug plot
  if (rugplot){
    rug(mod_frame[[z]])
  }
  # Add Histogram (Histogram only plots when minimum and maximum are the min/max of the moderator)
  #if (histogram & minimum=="min" & maximum=="max"){ ## DEACTIVATED
  if (histogram){
    par(new=T)
    hist(mod_frame[[z]], axes=F, xlab="", ylab="",main="", border=hist_col, col=hist_col)
    
    #   par(new=T)
    #    hist(subset(mod_frame[[z]], mod_frame[[w]]==1), axes=F, xlab="", ylab="",main="", ylim=c(0,1), border=hist_col, col=adjustcolor("red", alpha.f = 0.10))
    #    par(new=T)
    #    hist(subset(mod_frame[[z]], mod_frame[[w]]==0), axes=F, xlab="", ylab="",main="", border=hist_col, col=adjustcolor("blue", alpha.f = 0.10))
    
  }
}








interaction_plot_continuous <- function(model, effect, moderator, interaction, varcov="default", minimum="min", maximum="max", incr="default", num_points = 50, conf=.95, mean=FALSE, median=FALSE, alph=80, rugplot=T, histogram=T, title="Marginal effects plot", xlabel="Value of moderator", ylabel="Estimated marginal coefficient"){
  
  # Define a function to make colors transparent
  makeTransparent<-function(someColor, alpha=alph){
    newColor<-col2rgb(someColor)
    apply(newColor, 2, function(curcoldata){rgb(red=curcoldata[1], green=curcoldata[2],
                                                blue=curcoldata[3],alpha=alpha, maxColorValue=255)})
  }
  
  # Extract Variance Covariance matrix
  covMat = varcov
  
  # Extract the data frame of the model
  mod_frame = model.frame(model)
  
  # Get coefficients of variables
  beta_1 = model$coefficients[[effect]]
  beta_3 = model$coefficients[[interaction]]
  
  # Set range of the moderator variable
  # Minimum
  if (minimum == "min"){
    min_val = min(mod_frame[[moderator]])
  }else{
    min_val = minimum
  }
  # Maximum
  if (maximum == "max"){
    max_val = max(mod_frame[[moderator]])
  }else{
    max_val = maximum
  }
  
  # Check if minimum smaller than maximum
  if (min_val > max_val){
    stop("Error: Minimum moderator value greater than maximum value.")
  }
  
  # Determine intervals between values of the moderator
  if (incr == "default"){
    increment = (max_val - min_val)/(num_points - 1)
  }else{
    increment = incr
  }
  
  # Create list of moderator values at which marginal effect is evaluated
  x_2 <- seq(from=min_val, to=max_val, by=increment)
  
  # Compute marginal effects
  delta_1 = beta_1 + beta_3*x_2
  
  # Compute variances
  var_1 = covMat[effect,effect] + (x_2^2)*covMat[interaction, interaction] + 2*x_2*covMat[effect, interaction]
  
  # Standard errors
  se_1 = sqrt(var_1)
  
  # Upper and lower confidence bounds
  z_score = qnorm(1 - ((1 - conf)/2))
  upper_bound = delta_1 + z_score*se_1
  lower_bound = delta_1 - z_score*se_1
  
  # Determine the bounds of the graphing area
  max_y = max(upper_bound)
  min_y = min(lower_bound)
  
  # Make the histogram color
  hist_col = makeTransparent("grey")
  
  # Initialize plotting window
  plot(x=c(), y=c(), ylim=c(min_y, max_y), xlim=c(min_val, max_val), xlab=xlabel, ylab=ylabel, main=title, axes=F)
  axis(side=1)
  axis(side=2)
  
  # Plot estimated effects
  lines(y=delta_1, x=x_2, col="black", lty =1, lwd = 2)
  lines(y=upper_bound, x=x_2, lty=2)
  lines(y=lower_bound, x=x_2, lty=2)
  
  polygon(x = c(x_2, rev(x_2)),
          y = c(lower_bound, 
                rev(upper_bound)),
          col =  adjustcolor("red", alpha.f = 0.30), border = NA)
  
  # Add a dashed horizontal line for zero
  abline(h=0, lty=3)
  
  # Add a vertical line at the mean
  if (mean){
    abline(v = mean(mod_frame[[moderator]]), lty=2, col="red")
  }
  
  # Add a vertical line at the median
  if (median){
    abline(v = median(mod_frame[[moderator]]), lty=3, col="blue")
  }
  
  # Add Rug plot
  if (rugplot){
    rug(mod_frame[[moderator]])
  }
  # Add Histogram (Histogram only plots when minimum and maximum are the min/max of the moderator)
  #if (histogram & minimum=="min" & maximum=="max"){  ## DEACTIVATED 
  if (histogram){
    par(new=T)
    hist(mod_frame[[moderator]], axes=F, xlab="", ylab="",main="", border=hist_col, col=hist_col)
  }
}
